// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <cmath>
#include <limits>
#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;

class GreatestLeastTest : public functions::test::FunctionBaseTest {
 protected:
  template <typename T>
  void runTest(
      const std::string& query,
      const std::vector<std::vector<T>>& inputs,
      const std::vector<std::optional<T>>& output,
      std::optional<size_t> stringBuffersExpectedCount = std::nullopt,
      const TypePtr& type = CppToType<T>::create(),
      const TypePtr& resultType = CppToType<T>::create()) {
    // Create input vectors
    auto vectorSize = inputs[0].size();
    std::vector<VectorPtr> inputColumns(inputs.size());
    for (auto i = 0; i < inputColumns.size(); ++i) {
      inputColumns[i] = make_flat_vector<T>(inputs[i], type);
      for (auto j = 0; j < vectorSize; ++j) {
        inputColumns[i]->as_flat_vector<T>()->set(j, inputs[i][j]);
      }
    }

    // Call evaluate to run the query on the created input
    auto result = evaluate<SimpleVector<T>>(
        query, make_row_vector(inputColumns), std::nullopt, resultType);
    for (int32_t i = 0; i < vectorSize; ++i) {
      if (output[i].has_value()) {
        ASSERT_EQ(result->value_at(i), output[i]);
      } else {
        ASSERT_TRUE(result->is_null_at(i));
      }
    }

    if (stringBuffersExpectedCount.has_value()) {
      ASSERT_EQ(
          *stringBuffersExpectedCount,
          result->template as_flat_vector<StringView>()->stringBuffers().size());
    }
  }

  void runDecimalTest(
      const std::string& query,
      const std::vector<VectorPtr>& input,
      const VectorPtr& output) {
    auto result = evaluate(query, make_row_vector(input));
    test::assertEqualVectors(output, result);
  }
};

TEST_F(GreatestLeastTest, leastDouble) {
  runTest<double>("least(c0)", {{0, 1.1, -1.1}}, {0, 1.1, -1.1});
  runTest<double>("least(c0, 1.0)", {{0, 1.1, -1.1}}, {0, 1, -1.1});
  runTest<double>(
      "least(c0, 1.0, c1)", {{0, 1.1, -1.1}, {100, -100, 0}}, {0, -100, -1.1});
}

TEST_F(GreatestLeastTest, leastReal) {
  runTest<float>("least(c0)", {{0, 1.1, -1.1}}, {0, 1.1, -1.1});
  runTest<float>("least(c0, 1.0::real)", {{0, 1.1, -1.1}}, {0, 1, -1.1});
  runTest<float>(
      "least(c0, 1.0::real, c1)",
      {{0, 1.1, -1.1}, {100, -100, 0}},
      {0, -100, -1.1});
}

TEST_F(GreatestLeastTest, greatestNanInput) {
  auto constexpr kInf32 = std::numeric_limits<float>::infinity();
  auto constexpr kInf64 = std::numeric_limits<double>::infinity();

  auto greatestFloat = [&](float a, float b, float c) {
    return evaluateOnce<float, float, float, float>(
               "greatest(c0, c1, c2)", {a}, {b}, {c})
        .value();
  };

  auto greatestDouble = [&](double a, double b, double c) {
    return evaluateOnce<double, double, double, double>(
               "greatest(c0, c1, c2)", {a}, {b}, {c})
        .value();
  };

  EXPECT_TRUE(std::isnan(greatestFloat(1.0, std::nanf("1"), 2.0)));
  EXPECT_TRUE(std::isnan(greatestFloat(std::nanf("1"), 1.0, kInf32)));

  EXPECT_TRUE(std::isnan(greatestDouble(1.0, std::nan("1"), 2.0)));
  EXPECT_TRUE(std::isnan(greatestDouble(std::nan("1"), 1.0, kInf64)));
}

TEST_F(GreatestLeastTest, leastNanInput) {
  auto constexpr kInf32 = std::numeric_limits<float>::infinity();
  auto constexpr kInf64 = std::numeric_limits<double>::infinity();

  auto leastFloat = [&](float a, float b, float c) {
    return evaluateOnce<float, float, float, float>(
               "least(c0, c1, c2)", {a}, {b}, {c})
        .value();
  };

  auto leastDouble = [&](double a, double b, double c) {
    return evaluateOnce<double, double, double, double>(
               "least(c0, c1, c2)", {a}, {b}, {c})
        .value();
  };

  EXPECT_EQ(leastFloat(1.0, std::nanf("1"), 0.5), 0.5);
  EXPECT_EQ(leastFloat(std::nanf("1"), 1.0, -kInf32), -kInf32);

  EXPECT_EQ(leastDouble(1.0, std::nan("1"), 0.5), 0.5);
  EXPECT_EQ(leastDouble(std::nan("1"), 1.0, -kInf64), -kInf64);
}

TEST_F(GreatestLeastTest, greatestDouble) {
  runTest<double>("greatest(c0)", {{0, 1.1, -1.1}}, {0, 1.1, -1.1});
  runTest<double>("greatest(c0, 1.0)", {{0, 1.1, -1.1}}, {1, 1.1, 1});
  runTest<double>(
      "greatest(c0, 1.0, c1)", {{0, 1.1, -1.1}, {100, -100, 0}}, {100, 1.1, 1});
}

TEST_F(GreatestLeastTest, greatestReal) {
  runTest<float>("greatest(c0)", {{0, 1.1, -1.1}}, {0, 1.1, -1.1});
  runTest<float>("greatest(c0, 1.0::real)", {{0, 1.1, -1.1}}, {1, 1.1, 1});
  runTest<float>(
      "greatest(c0, 1.0::real, c1)",
      {{0, 1.1, -1.1}, {100, -100, 0}},
      {100, 1.1, 1});
}

TEST_F(GreatestLeastTest, leastInteger) {
  // TinyInt
  runTest<int8_t>("least(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int8_t>("least(c0, c1)", {{0, 1, -1}, {100, -100, 0}}, {0, -100, -1});
  // SmallInt
  runTest<int16_t>("least(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int16_t>(
      "least(c0, c1)", {{0, 1, -1}, {100, -100, 0}}, {0, -100, -1});
  // Integer
  runTest<int32_t>("least(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int32_t>(
      "least(c0, c1)", {{0, 1, -1}, {100, -100, 0}}, {0, -100, -1});
  // BigInt
  runTest<int64_t>("least(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int64_t>("least(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int64_t>(
      "least(c0, c1)", {{0, 1, -1}, {100, -100, 0}}, {0, -100, -1});
  runTest<int64_t>("least(c0, 1)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int64_t>(
      "least(c0, 1 , c1)", {{0, 1, -1}, {100, -100, 0}}, {0, -100, -1});
}

TEST_F(GreatestLeastTest, greatestInteger) {
  // TinyInt
  runTest<int8_t>("greatest(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int8_t>(
      "greatest(c0, c1)", {{0, 1, -1}, {100, -100, 0}}, {100, 1, 0});
  // SmallInt
  runTest<int16_t>("greatest(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int16_t>(
      "greatest(c0, c1)", {{0, 1, -1}, {100, -100, 0}}, {100, 1, 0});
  // Integer
  runTest<int32_t>("greatest(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int32_t>(
      "greatest(c0, c1)", {{0, 1, -1}, {100, -100, 0}}, {100, 1, 0});
  // BigInt
  runTest<int64_t>("greatest(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int64_t>("greatest(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int64_t>("greatest(c0)", {{0, 1, -1}}, {0, 1, -1});
  runTest<int64_t>("greatest(c0, 1)", {{0, 1, -1}}, {1, 1, 1});
  runTest<int64_t>(
      "greatest(c0, 1 , c1)", {{0, 1, -1}, {100, -100, 0}}, {100, 1, 1});
}

TEST_F(GreatestLeastTest, greatestVarchar) {
  runTest<StringView>(
      "greatest(c0)", {{"a"_sv, "b"_sv, "c"_sv}}, {"a"_sv, "b"_sv, "c"_sv});

  runTest<StringView>(
      "greatest(c0, 'aaa')", {{""_sv, "abb"_sv}}, {"aaa"_sv, "abb"_sv});
}

TEST_F(GreatestLeastTest, leasstVarchar) {
  runTest<StringView>(
      "least(c0)", {{"a"_sv, "b"_sv, "c"_sv}}, {"a"_sv, "b"_sv, "c"_sv});

  runTest<StringView>(
      "least(c0, 'aaa')", {{""_sv, "abb"_sv}}, {""_sv, "aaa"_sv});
}

TEST_F(GreatestLeastTest, greatestTimeStamp) {
  runTest<Timestamp>(
      "greatest(c0, c1, c2)",
      {{Timestamp(0, 0), Timestamp(10, 100), Timestamp(100, 10)},
       {Timestamp(1, 0), Timestamp(10, 1), Timestamp(100, 10)},
       {Timestamp(0, 1), Timestamp(312, 100), Timestamp(100, 11)}},
      {Timestamp(1, 0), Timestamp(312, 100), Timestamp(100, 11)});
}

TEST_F(GreatestLeastTest, leastTimeStamp) {
  runTest<Timestamp>(
      "least(c0, c1, c2)",
      {{Timestamp(0, 0), Timestamp(10, 100), Timestamp(100, 10)},
       {Timestamp(1, 0), Timestamp(10, 1), Timestamp(100, 10)},
       {Timestamp(0, 1), Timestamp(312, 100), Timestamp(1, 10)}},
      {Timestamp(0, 0), Timestamp(10, 1), Timestamp(1, 10)});
}

TEST_F(GreatestLeastTest, greatestTimestampWithTimezone) {
  auto greatest = [&](const std::string& a,
                      const std::string& b,
                      const std::string& c) {
    auto result = evaluateOnce<std::string>(
        "cast(greatest(cast(c0 as timestamp with time zone), cast(c1 as timestamp with time zone), cast(c2 as timestamp with time zone)) as varchar)",
        std::optional(a),
        std::optional(b),
        std::optional(c));
    return result.value();
  };

  auto least = [&](const std::string& a,
                   const std::string& b,
                   const std::string& c) {
    auto result = evaluateOnce<std::string>(
        "cast(least(cast(c0 as timestamp with time zone), cast(c1 as timestamp with time zone), cast(c2 as timestamp with time zone)) as varchar)",
        std::optional(a),
        std::optional(b),
        std::optional(c));
    return result.value();
  };

  EXPECT_EQ(
      "2024-04-10 08:11:22.010 America/Los_Angeles",
      greatest(
          "2024-04-10 10:11:22.01 America/New_York",
          "2024-02-10 10:11:22.01 America/New_York",
          "2024-04-10 08:11:22.01 America/Los_Angeles"));
  EXPECT_EQ(
      "2024-02-10 10:11:22.010 America/New_York",
      least(
          "2024-04-10 10:11:22.01 America/New_York",
          "2024-02-10 10:11:22.01 America/New_York",
          "2024-04-10 08:11:22.01 America/Los_Angeles"));
}

TEST_F(GreatestLeastTest, greatestDate) {
  runTest<int32_t>(
      "greatest(c0, c1, c2)",
      {
          {0, 5, 0},
          {1, 0, -5},
          {5, -5, -10},
      },
      {5, 5, 0},
      std::nullopt,
      DATE(),
      DATE());
}

TEST_F(GreatestLeastTest, leastDate) {
  runTest<int32_t>(
      "least(c0, c1, c2)",
      {{0, 0, 5}, {1, -1, -1}, {5, 5, -5}},
      {0, -1, -5},
      std::nullopt,
      DATE(),
      DATE());
}

TEST_F(GreatestLeastTest, greatestLeastIpAddress) {
  auto greatest = [&](const std::optional<std::string>& a,
                      const std::optional<std::string>& b,
                      const std::optional<std::string>& c) {
    return evaluateOnce<std::string>(
        "cast(greatest(cast(c0 as ipaddress), cast(c1 as ipaddress), cast(c2 as ipaddress)) as varchar)",
        a,
        b,
        c);
  };

  auto least = [&](const std::optional<std::string>& a,
                   const std::optional<std::string>& b,
                   const std::optional<std::string>& c) {
    return evaluateOnce<std::string>(
        "cast(least(cast(c0 as ipaddress), cast(c1 as ipaddress), cast(c2 as ipaddress)) as varchar)",
        a,
        b,
        c);
  };

  auto greatestValue = greatest(
      "1.1.1.1", "255.255.255.255", "2001:0db8:0000:0000:0000:ff00:0042:832");
  EXPECT_EQ("2001:db8::ff00:42:832", greatestValue.value());

  auto leastValue = least(
      "1.1.1.1", "255.255.255.255", "2001:0db8:0000:0000:0000:ff00:0042:832");
  EXPECT_EQ("1.1.1.1", leastValue.value());

  auto greatestValueWithNulls =
      greatest("1.1.1.1", "255.255.255.255", std::nullopt);
  EXPECT_FALSE(greatestValueWithNulls.has_value());

  auto leastValueWithNulls = least(
      std::nullopt,
      "255.255.255.255",
      "2001:0db8:0000:0000:0000:ff00:0042:832");
  EXPECT_FALSE(leastValueWithNulls.has_value());
}

TEST_F(GreatestLeastTest, stringBuffersMoved) {
  runTest<StringView>(
      "least(c0, c1)",
      {{"aaaaaaaaaaaaaa"_sv, "bbbbbbbbbbbbbb"_sv},
       {"cccccccccccccc"_sv, "dddddddddddddd"_sv}},
      {"aaaaaaaaaaaaaa"_sv, "bbbbbbbbbbbbbb"_sv},
      1);

  runTest<StringView>(
      "least(c0, c1, '')",
      {{"aaaaaaaaaaaaaa"_sv, "bbbbbbbbbbbbbb"_sv},
       {"cccccccccccccc"_sv, "dddddddddddddd"_sv}},
      {""_sv, ""_sv},
      0);
}

TEST_F(GreatestLeastTest, clear_nulls) {
  auto c0 = make_flat_vector<int64_t>(10, melon::identity);
  auto result = evaluate<SimpleVector<int64_t>>(
      "SWITCH(c0 > 5, null::BIGINT, greatest(c0))", make_row_vector({c0}));
  ASSERT_EQ(result->size(), 10);
  for (int i = 0; i < result->size(); ++i) {
    if (i > 5) {
      ASSERT_TRUE(result->is_null_at(i));
    } else {
      ASSERT_EQ(result->value_at(i), i);
    }
  }
}

TEST_F(GreatestLeastTest, shortDecimal) {
  const auto type = DECIMAL(10, 4);
  static const auto kMin = DecimalUtil::kLongDecimalMin + 1;
  static const auto kMax = DecimalUtil::kLongDecimalMax - 1;

  const auto a = make_nullable_flat_vector<int64_t>(
      {10000, -10000, 20000, kMax, kMin, std::nullopt}, type);
  const auto b = make_nullable_flat_vector<int64_t>(
      {-10000, 10000, -20000, kMin, kMax, 1}, type);
  runDecimalTest("least(c0)", {a}, a);
  runDecimalTest("greatest(c0)", {a}, a);

  auto expected = make_nullable_flat_vector<int64_t>(
      {-10000, -10000, -20000, kMin, kMin, std::nullopt}, type);
  runDecimalTest("least(c0, c1)", {a, b}, expected);

  expected = make_nullable_flat_vector<int64_t>(
      {10000, 10000, 20000, kMax, kMax, std::nullopt}, type);
  runDecimalTest("greatest(c0, c1)", {a, b}, expected);
}

TEST_F(GreatestLeastTest, longDecimal) {
  const auto type = DECIMAL(38, 10);
  static const auto kMin = DecimalUtil::kLongDecimalMin + 1;
  static const auto kMax = DecimalUtil::kLongDecimalMax - 1;

  const auto a = make_nullable_flat_vector<int128_t>(
      {HugeInt::build(10, 300),
       HugeInt::build(-10, 300),
       HugeInt::build(200, 300),
       kMax,
       kMin,
       std::nullopt},
      type);
  const auto b = make_nullable_flat_vector<int128_t>(
      {HugeInt::build(-10, 300),
       HugeInt::build(10, 300),
       HugeInt::build(-200, 300),
       kMin,
       kMax,
       HugeInt::build(1, 1)},
      type);
  runDecimalTest("least(c0)", {a}, a);
  runDecimalTest("greatest(c0)", {a}, a);

  auto expected = make_nullable_flat_vector<int128_t>(
      {HugeInt::build(-10, 300),
       HugeInt::build(-10, 300),
       HugeInt::build(-200, 300),
       kMin,
       kMin,
       std::nullopt},
      type);
  runDecimalTest("least(c0, c1)", {a, b}, expected);

  expected = make_nullable_flat_vector<int128_t>(
      {HugeInt::build(10, 300),
       HugeInt::build(10, 300),
       HugeInt::build(200, 300),
       kMax,
       kMax,
       std::nullopt},
      type);
  runDecimalTest("greatest(c0, c1)", {a, b}, expected);
}

TEST_F(GreatestLeastTest, boolean) {
  auto data = make_row_vector({
      make_flat_vector<bool>({true, true, false, false, true, false}),
      make_nullable_flat_vector<bool>(
          {true, false, true, false, std::nullopt, std::nullopt}),
  });

  auto result = evaluate("least(c0, c1)", data);
  auto expected = make_nullable_flat_vector<bool>(
      {true, false, false, false, std::nullopt, std::nullopt});
  test::assertEqualVectors(expected, result);

  result = evaluate("greatest(c0, c1)", data);
  expected = make_nullable_flat_vector<bool>(
      {true, true, true, false, std::nullopt, std::nullopt});
  test::assertEqualVectors(expected, result);
}
